import numpy as np
import torch
import cma
import logging
from datetime import datetime
from pyted.frontend import Variable
from pyted import get_circuit, DataSweep, debug
from frontend.torch import TorchVariable
from TuRBOat import TurboAt


class BOES_Demo:
    def __init__(
        self,
        tester,
        problem,
        pop_size: int = None,
        maxiters: int = 100,
        seed: int = 10,
        tracer=None,
        ref_point=None,
    ):
        assert maxiters > 2, "Maxiters should be larger than 2"
        self.seed = seed
        self.problem = problem
        self.tester = tester
        Dim = len(get_circuit().variables)
        # self.pop_size = pop_size
        if pop_size is None:
            self.pop_size = int(4 + 3 * np.log(Dim))
        self.maxiters = maxiters
        if ref_point is not None:
            self.ref_point = ref_point
        else:
            with get_circuit().setenv(quite=True, shell=False):
                res = self.tester.build()
            prob = self.problem(res)
            self._obj = prob.objective[0]
            obj = self._obj.clone().detach().numpy()
            self.ref_point = abs(obj) if obj!=0 else 1
        self.tracer = tracer(self.ref_point)
        self.params = [str(v) for v in get_circuit().variables.values()]
        self.init = [v.eval() for v in get_circuit().variables.values()]
        bound = [(v._lb, v._ub) for _, v in get_circuit().variables.items()]
        self.lbs, self.ubs = map(list, zip(*bound))
        # Considering tech
        pdk_lmin_dict = dict(t40=0.04, t65=0.065, t180=0.18)
        pdk_name = get_circuit().get_pdk().name
        self.lmin = pdk_lmin_dict[pdk_name]
        self.vname = [name for name, _ in get_circuit().variables.items()]
        for i in range(Dim):
            if self.vname[i].lower().startswith("l") or self.vname[i].lower().endswith("l"): 
                self.lbs[i] = max(self.lbs[i],self.lmin)
                self.ubs[i] = min(self.ubs[i],50*self.lmin) 
            if self.vname[i].lower().startswith("w") or self.vname[i].lower().startswith("s")\
               or self.vname[i].lower().endswith("w") or self.vname[i].lower().endswith("s"):
                self.lbs[i] = max(self.lbs[i],4*self.lmin)
                self.ubs[i] = min(self.ubs[i],100*self.lmin) 
        # Redefine the search space according to the init
        for i in range(Dim):
            if self.init[i] < self.lbs[i]:
                self.lbs[i] = self.init[i]
            if self.init[i] > self.ubs[i]:
                self.ubs[i] = self.init[i]
        self.x0 = [(i-l)/(u-l) for i,l,u in zip(self.init,self.lbs,self.ubs)]
        #with open("opt","w") as handler:
            #for name, l, u , i in zip(self.vname,self.lbs,self.ubs,self.init):
                #print(name,"\t",l,"\t",u,"\t",i,file=handler)

        self.turboAt = TurboAt(
            lb=np.zeros(Dim),
            ub=np.ones(Dim),
            n_init=2*self.pop_size,
            #max_evals=2*self.pop_size+self.pop_size*(self.maxiters-2),
            max_evals=np.inf,
            init=np.atleast_2d(self.x0),
            with_init=True,
            batch_size=self.pop_size,
            verbose=True,
            use_ard=True,
            max_cholesky_size=2000,
            n_training_steps=50,
            min_cuda=1024,
            device="cpu",
            dtype="float64",
        )

        self.es = cma.CMAEvolutionStrategy(
            self.x0,
            0.1,
            {'popsize': self.pop_size,
             'maxiter': self.maxiters,
             'bounds': [0, 1],
             'seed': self.seed,
             'verbose': 3}
        )

        # def obj_func(x):
        #     # # WRITE CODE HERE within two '#' bar                                        #
        #     # # Converts the input x to list form and dim(x)=1                            #
        #     # #############################################################################

        #     # #############################################################################
        #     tester = self.update(x)
        #     prob = self.problem(tester)

        #     self._obj = prob.objective[0]
        #     obj = self._obj.clone().detach().numpy()
        #     return obj

    def obj_func(self, x):
        OV = []
        CV = []
        for i in range(len(x)):
            tester = self.update(x[i])
            prob = self.problem(tester)
            self._obj = prob.objective[0]
            obj = self._obj.clone().detach().numpy()
            cv = [x.tolist() for x in prob.constraint]
            OV.append(obj)
            CV.append(cv)

            # 记录当前优化迭代信息
            if self.tracer != None:
                self.tracer.collect_n_iters_fom(
                    prob.objective[0].clone().detach().numpy(),
                    np.array([con.to_numpy() for con in prob.constraint]),
                )

        CV = np.vstack(CV)
        penalty_ueq = np.maximum(CV, 0).sum(axis=1)

        foms = np.array(OV)/self.ref_point + 1e3*penalty_ueq
        return foms.tolist()

    def parallel_func(self, new_x):
        keylib = []
        reslib = []
        newlib = []
        get_circuit().datasweep = DataSweep("sweepdata", self.params, new_x.T)
        with get_circuit().setenv(shell=False):
            res = self.tester.build()
            # debug(res)
            for k, v in res.items():
                keylib.append(k)
                #debug("v", v)
                #debug("reslib", reslib)
                if isinstance(v,int):
                    v = [v]*len(new_x)
                if isinstance(v[0],int):
                    v = [TorchVariable(torch.tensor(i)) for i in v]
                reslib.append(v)
            reslib = torch.Tensor(reslib).T
            if reslib.ndim == 1:
                reslib = reslib.view(1, -1)
            for j in range(len(reslib)):
                newdict = {}
                for index, value in enumerate(keylib):
                    newdict[value] = TorchVariable(reslib[j][index])
                newlib.append(newdict)
        tests = [self.problem(i) for i in newlib]
        OV = [i.objective[0] for i in tests]
        CV = [[list(i.constraint)] for i in tests]

        if self.tracer != None:
            for obj, cv in zip(OV, CV):
                self.tracer.collect_n_iters_fom(
                    obj.clone().detach().numpy(),
                    np.array([con.to_numpy() for con in cv[0]]),
                )

        OV = np.array(OV)
        CV = np.vstack(CV)
        penalty_ueq = np.maximum(CV, 0).sum(axis=1)

        foms = OV/self.ref_point + 10*penalty_ueq
        foms = foms[:,np.newaxis]
        # debug("foms", foms.shape)
        return foms, OV, CV 
        # ret_values = []
        # ret_keys = list(newlib[0].keys())
        # for x in newlib:
        #     ret_values.append(list(x.values()))
        # debug("ret_keys", ret_keys)
        # debug("ret_values", ret_values)
        # return ret_keys, ret_values

    def solve(self, prophetVars=None):
        """
        prophetVars: 先验解(可选)
        """
        #logger = logging.getLogger("BOES")
        #logger.setLevel(logging.DEBUG)
        #f_handler = logging.FileHandler("BOES.log")
        #f_handler.setLevel(logging.DEBUG)
        #formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(lineno)s\n%(message)s\n")
        #f_handler.setFormatter(formatter)
        #logger.addHandler(f_handler)
        #logger.info(
        #    "Optimizer BOES enters the solve function\n"+\
        #    "The pop_size is {}\n".format(self.pop_size)+\
        #    "The maxiters is {}\n".format(self.maxiters)+\
        #    "The ref_point is {}\n".format(self.ref_point)
        #)
        start_time = datetime.now()
        with open("BOES.log","w") as handler:
            print(f"{start_time} - INFO - ", file=handler)
            print(
                f"\tOptimizer BOES enters the solve function\n"+\
                f"\tThe pop_size is {self.pop_size}\n"+\
                f"\tThe maxiters is {self.maxiters}\n"+\
                f"\tThe ref_point is {self.ref_point}\n", 
                f"\tThe seed is {self.seed}", 
                file=handler
            )
        seed = self.seed
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        best_x = None
        best_y = np.inf       
        best_o = None
        best_c = None
        best_iter = 0
        best_index = 0
        best_algo = None
        iter_cnt = 0
        #while not self.turboAt.stop():
        while iter_cnt < self.maxiters:
            iter_cnt += 1
            x_turbo = self.turboAt.ask(self.pop_size).tolist()
            x_cma = self.es.ask(self.pop_size)
            x = x_turbo + x_cma
            solutions = np.array(x)*(np.array(self.ubs)-np.array(self.lbs)) + \
                np.array(self.lbs)
            #with open("record","a") as handler:
                #print("x",x,file=handler)
                #print("s",solutions,file=handler)
            # solutions = solutions.tolist()
            foms, ov, cv = self.parallel_func(solutions)
            self.turboAt.tell(x_turbo, foms[:self.pop_size])
            self.es.tell(x_cma, foms[self.pop_size:])
            index = 0
            for temp_x, temp_y, temp_o, temp_c in zip(solutions, foms, ov, cv):
                if temp_y < best_y:
                    best_y = temp_y
                    best_x = temp_x
                    best_o = temp_o
                    best_c = temp_c.tolist()
                    best_iter = iter_cnt
                    best_index = index
                    if best_index < self.pop_size:
                        best_algo = "BO"
                    else:
                        best_algo = "ES"
                    temp_time = datetime.now()
                    delta_seconds = temp_time.timestamp() - start_time.timestamp()
                    delta_hms = self.seconds_to_hms(delta_seconds)
                    with open("Best.report","w") as handler:
                        print(f"Generated Time: {temp_time}", file=handler)
                        print(f"Elapsed Time Since Optimization Started: {delta_hms}", file=handler)
                        print("\nDesign Paramerers:", file=handler)
                        for name, val in zip(self.vname, best_x):
                            print(f"\t{name}:\t{val}", file=handler)
                        print("\nPerformances:", file=handler)
                        print(f"\tObjective:\t{best_o}", file=handler)
                        print(f"\tConstraints:\t{best_c[0]}", file=handler)
                        for c in best_c[0:]:
                            print(f"\t{' '*len('Constraints:')}\t{c}", file=handler)
                    with open("BOES.log","a") as handler:
                        print(f"{temp_time} - INFO - ", file=handler)
                        print(
                        f"\tIteration {iter_cnt}: find better candidate by {best_algo}!"+\
                        f" See Best.report for more details",
                        file=handler
                        )
                index += 1
        # best_x, best_y = self.de.run()
        iters_info = (
            (self.tracer.n_iters, self.tracer.n_iters_fom)
            if self.tracer != None
            else (None, None)
        )
        end_time = datetime.now()
        with open("BOES.log","a") as handler:
            print(f"{end_time} - INFO - ", file=handler)
            print(
            f"\tEnd optimizing, the best candidate is found in iteration {best_iter} by {best_algo}!"+\
            f" See Best.report for more details",
            file=handler
            )
        return {
            "success": True,
            "Vars": [best_x],
            "fom": iters_info,
        }

    def update(self, x):
        Variable.update_all_variables(x)
        with get_circuit().setenv(quite=True, shell=False):
            return self.tester.build()

    def seconds_to_hms(self, seconds):
        hours = int(seconds // 3600)
        minutes =int((seconds % 3600) // 60)
        seconds = int(seconds % 60)
        return f"{hours:02d}:{minutes:02d}:{seconds:02d}"
